import json
import numpy as np
from ultralytics import YOLO
import os
"""
使用yolo训练好的猪体尺模型标记猪的体尺
"""
# Load a model
model = YOLO(
    "C:/Users/kang_/Desktop/workspace/pig_pose/runs/pose/train3/weights/best.pt"
)  # load a pretrained model (recommended for training)

# 需要标记的图片的目录根
# 会自动便利子文件夹的 1280x720_MA.json 和 1280x720_MA.png文件
# 自动预测并将体尺数据输出到 1280x720_MA.json 中
# 因为会破坏数据, 请在使用前备份原来的标记数据
root = "C:/Users/kang_/Desktop/data/240927_pm"

for root, dirs, files in os.walk(root):
    for file in files:
        if not file.endswith(".json"):
            continue
        json_file_path = os.path.join(root, "1280x720_MA.json")
        with open(json_file_path, mode="r") as fp:
            label_json = json.load(fp)
        img_path = os.path.join(root, "1280x720_MA.png")
        rst = model.predict(source=img_path, conf=0.5)
        shapes = label_json["shapes"]
        for r in rst:
            xy_list = r.keypoints.xy[0]
            group_id = 0
            for xy in np.array(xy_list).tolist():
                pig_pose =                     {
                        "label": "pig_pose",
                        "points": [xy],
                        "group_id": group_id,
                        "description": "label_by_yolo",
                        "shape_type": "point",
                        "flags": {},
                        "mask": None,
                    }
                shapes.append(pig_pose)
                group_id += 1
        with open(file=json_file_path, mode="w") as fp:
            json.dump(label_json, fp)
